#!/usr/bin/python
import time
import os
from numpy import *
from kalmanFuncs import *


import csv
import h5py


# This file has a simulated input into the Kalman filter in either the 7 or 13 state implementation
# Currently, the angular rates can only be constant


# Estimate the 9 sensor biases (13) or 3 biases (7)
numberStates = 13

# Environmental constants
# g = 9.81
##### change g to a vector here
g = [0, 0, -9.81]

# hloc = array([[0.225],[-0.057],[-0.4378]])       ##### shizheli huanjing cichang shu
hloc = array([[1], [1], [1]])

# Sensor noise constants
# hnoise = array([0.01,0.01,0.01])
# anoise = array([0.01,0.01,0.01])
hnoise = array([0.001,0.001,0.001])
anoise = array([0.001,0.001,0.001])

# Process noise constants
# hnoisep = array([0.001,0.001,0.001])
# anoisep = array([0.001,0.001,0.001])
# gnoisep = array([0.01,0.01,0.01])
# wnoise = 0.002
hnoisep = array([0.0001,0.0001,0.0001])
anoisep = array([0.001,0.001,0.001])
gnoisep = array([0.001,0.001,0.001])
wnoise = 0.0002
# hnoisep = array([0.01,0.01,0.01])
# anoisep = array([0.01,0.01,0.01])
# gnoisep = array([0.1,0.1,0.1])
# wnoise = 0.02

# Simulation constants
t = 0.0
filterFreq = 200.0 #hz
##### filterFreq zheli yaogai, yinggai shi 200Hz

dt = 1.0 / filterFreq
outputFreq = 1.0 #hz
counter = filterFreq / outputFreq
e1dot = 0.05
e2dot = 0.03
e3dot = 0.02

if numberStates == 7:
	hnoisep *= 0
	anoisep *= 0

# Simulation Noise & Bias
hnoises = array([0.01,0.01,0.01])
anoises = array([0.01,0.01,0.01])
gnoises = array([0.01,0.01,0.01])
hbiass = array([0.01,0.01,0.01])
abiass = array([0.01,0.01,0.01])
gbiass = array([0.05,0.05,0.05])

# Initial Conditions
# X = array([[1],[0],[0],[0],[0],[0],[0],[0],[0],[0],[0],[0],[0]])

#### P zheli buzhidao xubux xuyao gai
P = diag([0,0,0,0,0,0,0,0,0,0,0,0,0])
# uu = array([[e1dot],[e2dot],[e3dot],[0],[0],[0],[0],[0],[0]])
uu = array([[0],[0],[0],[0],[0],[0],[0],[0],[0]])


dataCnt = 1
fileOut = open('EKFresult.csv', 'w')
csv_writer = csv.writer(fileOut)
csv_writer.writerow(["q0", "q1", "q2", "q3", "wp", "wq", "wr", "ax", "ay", "az"])

#### load data file
path = "/home/ubuntu/user_space/Ronin/root_dir/a028_3/data.hdf5"
f = h5py.File(path, "a")

raw = f["synced"]
gyro_s = raw["gyro"]
acce = raw["acce"]
###### q and v initial value to be set
px_init = 0.32157008
py_init = -0.10793696
pz_init = 0.00769295

vx_init = (0.32156168 - 0.32157008) / 0.005
vy_init = (-0.10786508 - -0.10793696) / 0.005
vz_init = (0.00769964 - 0.00769295) / 0.005
### use the p value for a
ax_init = ((0.32155362 - 0.32156168) / 0.005 - (0.32156168 - 0.32157008) / 0.005) / 0.005
ay_init = ((-0.10778917 - -0.10786508 ) / 0.005 - (-0.10786508 - -0.10793696) / 0.005) / 0.005
az_init = ((0.00770835 - 0.00769964) / 0.005 - (0.00769964 - 0.00769295) / 0.005)/ 0.005

##### q = 0.54602261  0.49708395  0.49287465 -0.46026271
# q0 = -0.51167799
# q1 = 0.55961401
# q2 = -0.44363793
# q3 = -0.47770652
q0 = 0.54602261
q1 = 0.49708395
q2 = 0.49287465
q3 = -0.46026271
##### w =  0.0297434  -0.01214644  0.00723207
X = array([[q0],[q1],[q2],[q3],[0.0297434],[-0.01214644],[0.00723207],[ax_init],[ay_init],[az_init],[0],[0],[0]])

# Main filter Loop
while (True & (dataCnt < 30000)):
	
	# Time update
	t += dt
	counter -= 1
	Q = Qmat(X,dt,wnoise,gnoisep,anoisep,hnoisep)
	A = Amat(X,uu,g,dt)
	X = Afunc(X,uu,g,dt)
	P = dot(dot(A,P),A.T) + Q
	H = Hmat(X,hloc,g)
	R = Rnmat(hnoise,anoise)
	K = dot(dot(P, H.T), linalg.inv(dot(dot(H, P), H.T) + R))
	YH = Hfunc(X,g,hloc)

	# path = "/home/ubuntu/user_space/Ronin/root_dir/a027_3/data.hdf5"
	# f = h5py.File(path, "a")
	#
	# raw = f["synced"]
	# gyro = raw["gyro"]
	# acce = raw["acce"]
	#### adjust input here
	
	# Measurement generation
	e1 = arctan(tan(e1dot * t / 2.)) * 2
	e2 = arctan(tan(e2dot * t / 2.)) * 2
	e3 = arctan(tan(e3dot * t / 2.)) * 2
	# q0 = cos(e1/2)*cos(e2/2)*cos(e3/2)+sin(e1/2)*sin(e2/2)*sin(e3/2)
	# q1 = sin(e1/2)*cos(e2/2)*cos(e3/2)-cos(e1/2)*sin(e2/2)*sin(e3/2)
	# q2 = cos(e1/2)*sin(e2/2)*cos(e3/2)+sin(e1/2)*cos(e2/2)*sin(e3/2)
	# q3 = cos(e1/2)*cos(e2/2)*sin(e3/2)-sin(e1/2)*sin(e2/2)*cos(e3/2)
	# q0 = X[0][0]
	# q1 = X[1][0]
	# q2 = X[2][0]
	# q3 = X[3][0]

	# q0 = -0.51167799
	# q1 = 0.55961401
	# q2 = -0.44363793
	# q3 = -0.47770652
	# w1 = e1dot * cos(e2) * cos(e3) + e2dot * sin(e3)
	# w2 = -e1dot * cos(e2) * sin(e3) + e2dot * cos(e3)
	# w3 = e1dot * sin(e2) + e3dot
	# w1 = gyro[dataCnt][3]
	# w2 = gyro[dataCnt][3]
	# w3 = gyro[dataCnt][3] wrong index here
	w1 = gyro_s[dataCnt][0]
	w2 = gyro_s[dataCnt][1]
	w3 = gyro_s[dataCnt][2]

	# print(acce.shape)
	# print(gyro.shape)


# euler angle determination of rotation matrix - seems to require a negative sign on calculation of angles from estimated quaternion?
#	Tbi = array([[cos(e2)*cos(e3),-cos(e2)*sin(e3),sin(e2)],[sin(e1)*sin(e2)*cos(e3)+sin(e3)*cos(e3),-sin(e1)*sin(e2)*sin(e3)+cos(e3)*cos(e1),-sin(e1)*cos(e2)],[-cos(e1)*sin(e2)*cos(e3)+sin(e3)*sin(e1),cos(e1)*sin(e2)*sin(e3)+cos(e3)*sin(e1),cos(e1)*cos(e2)]])

# quaternion determination of rotation matrix
	##### Tb zheli you typo
	# Tbi = array([[1-2*(q2**2+q3**2),2*(q1*q2+q0*q3),2*(q1*q3-q0*q2)],[2*(q1*q2-q0*q3),1-2*(q1**2+q3**2),2*(q2*q3+q0*q1)],[2*(q1*q3+q0*q2),2*(q2*q3-q0*q1),1-2*(q1**2+q2**2)]])
	Tbi = array([[1 - 2 * (q2 ** 2 + q3 ** 2), 2 * (q1 * q2 - q0 * q3), 2 * (q1 * q3 + q0 * q2)],
				 [2 * (q1 * q2 + q0 * q3), 1 - 2 * (q1 ** 2 + q3 ** 2), 2 * (q2 * q3 - q0 * q1)],
				 [2 * (q1 * q3 - q0 * q2), 2 * (q2 * q3 + q0 * q1), 1 - 2 * (q1 ** 2 + q2 ** 2)]])
	
	acc = array([acce[dataCnt][0],acce[dataCnt][1],acce[dataCnt][2]])
	print(acc)
	##### change the acc direction here to Global coordinate
	acc = np.dot(Tbi, acc)
	print(acc)
	# X[7] = acc[0]
	# X[8] = acc[1]
	# X[9] = acc[2]

	# acc[0] += random.normal(abiass[0],anoises[0])
	# acc[1] += random.normal(abiass[1],anoises[1])
	# acc[2] += random.normal(abiass[2],anoises[2])
	gyro = array([w1,w2,w3])
	# gyro[0] += random.normal(gbiass[0],gnoises[0])
	# gyro[1] += random.normal(gbiass[1],gnoises[1])
	# gyro[2] += random.normal(gbiass[2],gnoises[2])
	mag = dot(Tbi,hloc)
	# mag[0] += random.normal(hbiass[0],hnoises[0])
	# mag[1] += random.normal(hbiass[1],hnoises[1])
	# mag[2] += random.normal(hbiass[2],hnoises[2])   ##### zheli yinggai wei 0
	mag[0] = 0
	mag[1] = 0
	mag[2] = 0

	### print mag shape
	print(mag.shape)

	####### Y zheli zai jiancha yixia
	####### Y de guance zhi zheli xuyao gai, zhiqian meigai, jia jiao sudu
	Y = array([[w1],[w2],[w3],[acc[0]],[acc[1]],[acc[2]]])
	# uu = array([[gyro[0]],[gyro[1]],[gyro[2]],[acc[0]],[acc[1]],[acc[2]],[mag[0]],[mag[1]],[mag[2]]])
	
	# Measurement update
	X = X + dot(K,(Y-YH))
	P = dot((identity(13)-dot(K,H)),P)


	dataCnt += 1

	# Output generation
	if counter == 0:
		counter = filterFreq/outputFreq

	Eq0 = X[0][0]
	Eq1 = X[1][0]
	Eq2 = X[2][0]
	Eq3 = X[3][0]
	Ewpb = X[4][0]
	Ewqb = X[5][0]
	Ewrb = X[6][0]
	Eaxb = X[7][0]
	Eayb = X[8][0]
	Eazb = X[9][0]
	Ehxb = X[10][0]
	Ehyb = X[11][0]
	Ehzb = X[12][0]

	print(Eaxb, Eayb, Eazb)

	vx_init = vx_init + Eaxb * dt
	vy_init = vy_init + Eayb * dt
	vz_init = vz_init + Eazb * dt

	px_init = px_init + vx_init * dt + 0.5 * Eaxb * dt ** 2
	py_init = py_init + vy_init * dt + 0.5 * Eayb * dt ** 2
	pz_init = pz_init + vz_init * dt + 0.5 * Eazb * dt ** 2

	q0 = Eq0
	q1 = Eq1
	q2 = Eq2
	q3 = Eq3


	##### Euler angles generation
	Ee1 = arctan2(2 * (Eq0 * Eq1 + Eq2 * Eq3), 1 - 2 * (Eq1 ** 2 + Eq2 ** 2))
	Ee2 = arcsin(2 * (Eq0 * Eq2 - Eq3 * Eq1))
	Ee3 = arctan2(2 * (Eq0 * Eq3 + Eq1 * Eq2), 1 - 2 * (Eq2 ** 2 + Eq3 ** 2))

	print("Print first place")

	# os.system("clear")
	# print('time - ', t)
	# print('meas, calc, error', set_printoptions(precision=4, linewidth=120, suppress=True))
	# print(Y.T)
	# print(YH.T)
	# print(Y.T - YH.T)
	#
	#
	# print("Print second place")
	#
	# print('Euler angles')
	# print
	# 'est %.4f \t\t real %.4f \t\t error %.4f' % (Ee1 * 180 / pi, e1 * 180 / pi, Ee1 * 180 / pi - e1 * 180 / pi)
	# print
	# 'est %.4f \t\t real %.4f \t\t error %.4f' % (Ee2 * 180 / pi, e2 * 180 / pi, Ee2 * 180 / pi - e2 * 180 / pi)
	# print
	# 'est %.4f \t\t real %.4f \t\t error %.4f' % (Ee3 * 180 / pi, e3 * 180 / pi, Ee3 * 180 / pi - e3 * 180 / pi)
	# print
	# sqrt(Eq0 ** 2 + Eq1 ** 2 + Eq2 ** 2 + Eq3 ** 2)
	#
	#
	# #### output state to a file
	# csv_writer.writerow([Eq0, Eq1, Eq2, Eq3, px_init, py_init, pz_init, Eaxb, Eayb, Eazb])
	#
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eq0, q0, Eq0 - q0))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eq1, q1, Eq1 - q1))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eq2, q2, Eq2 - q2))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eq3, q3, Eq3 - q3))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ewpb, gbiass[0], Ewpb - gbiass[0]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ewqb, gbiass[1], Ewqb - gbiass[1]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ewrb, gbiass[2], Ewrb - gbiass[2]))
	# # if numberStates == 7:
	# # 	continue
	#
	# print("Print 3rd place")
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eaxb, abiass[0], Eaxb - abiass[0]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eayb, abiass[1], Eayb - abiass[1]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Eazb, abiass[2], Eazb - abiass[2]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ehxb, hbiass[0], Ehxb - hbiass[0]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ehyb, hbiass[1], Ehyb - hbiass[1]))
	# print('est %.4f \t\t real %.4f \t\t error %.4f' % (Ehzb, hbiass[2], Ehzb - hbiass[2]))

fileOut.close()



